//
// Created by hjhu on 2024/7/2.
//

#ifndef TEST_WAV2LIP_CPP_CONVBLOCK_H
#define TEST_WAV2LIP_CPP_CONVBLOCK_H

#include <torch/torch.h>

namespace coastal {

    using namespace torch;
    using namespace std;
    namespace nn = nn;

    class ConvBlockImpl : public nn::Module {

    private:
        nn::BatchNorm2d bn1, bn2, bn3;
        nn::Conv2d conv1, conv2, conv3;
        std::optional<nn::Sequential> downsample;

    public:

        ConvBlockImpl(int64_t in_planes, int64_t out_planes);

        torch::Tensor forward(torch::Tensor x);

        static nn::Conv2d
        conv3x3(int64_t in_planes, int64_t out_planes, int64_t stride = 1, int64_t padding = 1, bool bias = false) {
            return nn::Conv2d(
                    nn::Conv2dOptions(in_planes, out_planes, 3).stride(stride).padding(padding).bias(false));
        }
    };

    TORCH_MODULE(ConvBlock);
}


#endif //TEST_WAV2LIP_CPP_CONVBLOCK_H
